{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp tst_net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# #hide\n",
    "# # from nbdev.showdoc import *\n",
    "# from fastcore.test import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # export\n",
    "# import torch.nn as nn\n",
    "# from collections import OrderedDict\n",
    "# from model_constructor.layers import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # export\n",
    "# import torch,math,sys\n",
    "# import torch.utils.model_zoo as model_zoo\n",
    "# from functools import partial\n",
    "# # from fastai.torch_core import Module"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Xresnet from fastai v1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## helpers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # export\n",
    "# act_fn = nn.ReLU(inplace=True)\n",
    "\n",
    "# # in layers\n",
    "# class Flatten(Module):\n",
    "#     def forward(self, x): return x.view(x.size(0), -1)\n",
    "\n",
    "# def init_cnn(m):\n",
    "#     if getattr(m, 'bias', None) is not None: nn.init.constant_(m.bias, 0)\n",
    "#     if isinstance(m, (nn.Conv2d,nn.Linear)): nn.init.kaiming_normal_(m.weight)\n",
    "#     for l in m.children(): init_cnn(l)\n",
    "\n",
    "# def conv(ni, nf, ks=3, stride=1, bias=False):\n",
    "#     return nn.Conv2d(ni, nf, kernel_size=ks, stride=stride, padding=ks//2, bias=bias)\n",
    "# # in layers\n",
    "# def noop(x): return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## conv layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # export\n",
    "\n",
    "# def conv_layer(ni, nf, ks=3, stride=1, zero_bn=False, act=True):\n",
    "#     bn = nn.BatchNorm2d(nf)\n",
    "#     nn.init.constant_(bn.weight, 0. if zero_bn else 1.)\n",
    "#     layers = [conv(ni, nf, ks, stride=stride), bn]\n",
    "#     if act: layers.append(act_fn)\n",
    "#     return nn.Sequential(*layers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ResBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # v1\n",
    "# class ResBlock(Module):\n",
    "#     def __init__(self, expansion, ni, nh, stride=1):\n",
    "#         nf,ni = nh*expansion,ni*expansion\n",
    "#         layers  = [conv_layer(ni, nh, 3, stride=stride),\n",
    "#                    conv_layer(nh, nf, 3, zero_bn=True, act=False)\n",
    "#         ] if expansion == 1 else [\n",
    "#                    conv_layer(ni, nh, 1),\n",
    "#                    conv_layer(nh, nh, 3, stride=stride),\n",
    "#                    conv_layer(nh, nf, 1, zero_bn=True, act=False)\n",
    "#         ]\n",
    "#         self.convs = nn.Sequential(*layers)\n",
    "#         # TODO: check whether act=True works better\n",
    "#         self.idconv = noop if ni==nf else conv_layer(ni, nf, 1, act=False)\n",
    "#         self.pool = noop if stride==1 else nn.AvgPool2d(2, ceil_mode=True)\n",
    "\n",
    "#     def forward(self, x): return act_fn(self.convs(x) + self.idconv(self.pool(x)))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # v2\n",
    "# class ResBlock(Module):\n",
    "#     def __init__(self, expansion, ni, nh, stride=1, \n",
    "#                  conv_layer=conv_layer, act_fn=act_fn,\n",
    "#                  pool=nn.AvgPool2d(2, ceil_mode=True)):\n",
    "#         nf,ni = nh*expansion,ni*expansion\n",
    "#         layers  = [(f\"conv_0\", conv_layer(ni, nh, 3, stride=stride)),\n",
    "#                    (f\"conv_1\", conv_layer(nh, nf, 3, zero_bn=True, act=False))\n",
    "#         ] if expansion == 1 else [\n",
    "#                    (f\"conv_0\",conv_layer(ni, nh, 1)),\n",
    "#                    (f\"conv_1\",conv_layer(nh, nh, 3, stride=stride)),\n",
    "#                    (f\"conv_2\",conv_layer(nh, nf, 1, zero_bn=True, act=False))\n",
    "#         ]\n",
    "#         self.convs = nn.Sequential(OrderedDict(layers))\n",
    "#         self.pool = noop if stride==1 else pool\n",
    "#         self.idconv = noop if ni==nf else conv_layer(ni, nf, 1, act=False)\n",
    "#         self.act_fn =act_fn\n",
    "\n",
    "#     def forward(self, x): return self.act_fn(self.convs(x) + self.idconv(self.pool(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # export\n",
    "# class ResBlock(Module):\n",
    "#     def __init__(self, expansion, ni, nh, stride=1, \n",
    "#                  conv_layer=conv_layer, act_fn=act_fn,\n",
    "#                  pool=nn.AvgPool2d(2, ceil_mode=True), sa=False,sym=False):\n",
    "#         nf,ni = nh*expansion,ni*expansion\n",
    "#         layers  = [(f\"conv_0\", conv_layer(ni, nh, 3, stride=stride)),\n",
    "#                    (f\"conv_1\", conv_layer(nh, nf, 3, zero_bn=True, act=False))\n",
    "#         ] if expansion == 1 else [\n",
    "#                    (f\"conv_0\",conv_layer(ni, nh, 1)),\n",
    "#                    (f\"conv_1\",conv_layer(nh, nh, 3, stride=stride)),\n",
    "#                    (f\"conv_2\",conv_layer(nh, nf, 1, zero_bn=True, act=False))\n",
    "#         ]\n",
    "#         if sa: layers.append(('sa', SimpleSelfAttention(nf,ks=1,sym=sym)))\n",
    "#         self.convs = nn.Sequential(OrderedDict(layers))\n",
    "#         self.pool = noop if stride==1 else pool\n",
    "#         self.idconv = noop if ni==nf else conv_layer(ni, nf, 1, act=False)\n",
    "#         self.act_fn =act_fn\n",
    "\n",
    "#     def forward(self, x): return self.act_fn(self.convs(x) + self.idconv(self.pool(x)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ResBlock(1,64,64,sa=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ResBlock(1,64,64,)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## first vers XResNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # initial version\n",
    "# class XResNet(nn.Sequential):\n",
    "#     def __init__(self, expansion, layers, c_in=3, c_out=1000):\n",
    "#         stem = []\n",
    "#         sizes = [c_in,32,32,64]\n",
    "#         for i in range(3):\n",
    "#             stem.append(conv_layer(sizes[i], sizes[i+1], stride=2 if i==0 else 1))\n",
    "          \n",
    "#         block_szs = [64//expansion,64,128,256,512] +[256]*(len(layers)-4)\n",
    "#         blocks = [self._make_layer(expansion, block_szs[i], block_szs[i+1], l, 1 if i==0 else 2)\n",
    "#                   for i,l in enumerate(layers)]\n",
    "#         super().__init__(\n",
    "#             *stem,\n",
    "#             nn.MaxPool2d(kernel_size=3, stride=2, padding=1),\n",
    "#             *blocks,\n",
    "#             nn.AdaptiveAvgPool2d(1), Flatten(),\n",
    "#             nn.Linear(block_szs[-1]*expansion, c_out),\n",
    "#         )\n",
    "#         init_cnn(self)\n",
    "\n",
    "#     def _make_layer(self, expansion, ni, nf, blocks, stride):\n",
    "#         return nn.Sequential(\n",
    "#             *[ResBlock(expansion, ni if i==0 else nf, nf, stride if i==0 else 1)\n",
    "#               for i in range(blocks)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # v2\n",
    "# class XResNet(nn.Sequential):\n",
    "#     def __init__(self, expansion, layers, c_in=3, c_out=1000):\n",
    "#         stem = []\n",
    "#         sizes = [c_in,32,32,64]\n",
    "#         for i in range(3):\n",
    "#             stem.append(conv_layer(sizes[i], sizes[i+1], stride=2 if i==0 else 1))\n",
    "#              block_szs = [64//expansion,64,128,256,512] +[256]*(len(layers)-4)\n",
    "#         blocks = [self._make_layer(expansion, block_szs[i], block_szs[i+1], l, 1 if i==0 else 2)\n",
    "#                   for i,l in enumerate(layers)]\n",
    "#         super().__init__(OrderedDict([\n",
    "#             ('stem', nn.Sequential(*stem, nn.MaxPool2d(kernel_size=3, stride=2, padding=1))),\n",
    "#             ('body', nn.Sequential(*blocks)),\n",
    "#             ('head', nn.Sequential(nn.AdaptiveAvgPool2d(1), Flatten(), \n",
    "#              nn.Linear(block_szs[-1]*expansion, c_out)))\n",
    "#         ]))\n",
    "#         init_cnn(self)\n",
    "\n",
    "#     def _make_layer(self, expansion, ni, nf, blocks, stride):\n",
    "#         return nn.Sequential(\n",
    "#             *[ResBlock(expansion, ni if i==0 else nf, nf, stride if i==0 else 1)\n",
    "#               for i in range(blocks)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # v3\n",
    "# class XResNet():\n",
    "#     def __init__(self, expansion=1, layers=[2,2,2,2], c_in=3, c_out=1000):\n",
    "#         super().__init__()\n",
    "#         self.c_in, self.c_out,self.expansion = c_in,c_out,expansion\n",
    "#         self.stem = []\n",
    "#         sizes = [c_in,32,32,64]\n",
    "#         for i in range(3):\n",
    "#             self.stem.append(conv_layer(sizes[i], sizes[i+1], stride=2 if i==0 else 1))\n",
    "#         self.block_szs = [64//expansion,64,128,256,512] +[256]*(len(layers)-4)\n",
    "#         self.blocks = [self._make_layer(expansion, self.block_szs[i], self.block_szs[i+1], l, 1 if i==0 else 2)\n",
    "#                   for i,l in enumerate(layers)]\n",
    "        \n",
    "\n",
    "#     def _make_layer(self, expansion, ni, nf, blocks, stride):\n",
    "#         return nn.Sequential(\n",
    "#             *[ResBlock(expansion, ni if i==0 else nf, nf, stride if i==0 else 1)\n",
    "#               for i in range(blocks)])\n",
    "#     def __call__(self):\n",
    "#         model = nn.Sequential(OrderedDict([\n",
    "#             ('stem', nn.Sequential(*self.stem, nn.MaxPool2d(kernel_size=3, stride=2, padding=1))),\n",
    "#             ('body', nn.Sequential(*self.blocks)),\n",
    "#             ('head', nn.Sequential(nn.AdaptiveAvgPool2d(1), Flatten(), \n",
    "#              nn.Linear(self.block_szs[-1]*self.expansion, self.c_out)))\n",
    "#         ]))\n",
    "#         init_cnn(model)\n",
    "#         return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # v4\n",
    "# class XResNet():\n",
    "#     def __init__(self, expansion=1, layers=[2,2,2,2], c_in=3, c_out=1000):\n",
    "#         super().__init__()\n",
    "#         self.c_in, self.c_out,self.expansion = c_in,c_out,expansion\n",
    "#         self.stem = []\n",
    "#         sizes = [c_in,32,32,64]\n",
    "#         for i in range(3):\n",
    "#             self.stem.append(conv_layer(sizes[i], sizes[i+1], stride=2 if i==0 else 1))\n",
    "#         self.stem.append(nn.MaxPool2d(kernel_size=3, stride=2, padding=1))\n",
    "#         block_szs = [64//expansion,64,128,256,512] +[256]*(len(layers)-4)\n",
    "#         self.blocks = [self._make_layer(expansion, block_szs[i], block_szs[i+1], l, 1 if i==0 else 2)\n",
    "#                   for i,l in enumerate(layers)]\n",
    "#         self.head = [nn.AdaptiveAvgPool2d(1), Flatten(), \n",
    "#              nn.Linear(block_szs[-1]*self.expansion, self.c_out)]\n",
    "\n",
    "#     def _make_layer(self, expansion, ni, nf, blocks, stride):\n",
    "#         return nn.Sequential(\n",
    "#             *[ResBlock(expansion, ni if i==0 else nf, nf, stride if i==0 else 1)\n",
    "#               for i in range(blocks)])\n",
    "#     def __call__(self):\n",
    "#         model = nn.Sequential(OrderedDict([\n",
    "#             ('stem', nn.Sequential(*self.stem)),\n",
    "#             ('body', nn.Sequential(*self.blocks)),\n",
    "#             ('head', nn.Sequential(*self.head))\n",
    "#         ]))\n",
    "#         init_cnn(model)\n",
    "#         return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # v5\n",
    "# class XResNet():\n",
    "#     def __init__(self, expansion=1, layers=[2,2,2,2], c_in=3, c_out=1000):\n",
    "#         super().__init__()\n",
    "#         self.c_in, self.c_out,self.expansion,self.layers = c_in,c_out,expansion,layers\n",
    "#         self.stem_sizes = [c_in,32,32,64]\n",
    "#         self.block_szs = [64//expansion,64,128,256,512] +[256]*(len(layers)-4)\n",
    "#         self.block = ResBlock\n",
    "#         self.norm = nn.BatchNorm2d\n",
    "#         self.act_fn=nn.ReLU(inplace=True)\n",
    "# #         self._conv_layer = conv_layer\n",
    "    \n",
    "# #     @property\n",
    "# #     def conv_layer(self): return self._conv_layer\n",
    "\n",
    "#     def conv_layer(self, ni, nf, ks=3, stride=1, zero_bn=False, act=True):\n",
    "#         bn = self.norm(nf)\n",
    "#         nn.init.constant_(bn.weight, 0. if zero_bn else 1.)\n",
    "#         layers = [('conv', conv(ni, nf, ks, stride=stride)), ('norm', bn)]\n",
    "#         if act: layers.append(('act_fn', self.act_fn))\n",
    "#         return nn.Sequential(OrderedDict(layers))\n",
    "    \n",
    "#     @property\n",
    "#     def stem(self):\n",
    "#         return self._make_stem()\n",
    "#     @property\n",
    "#     def head(self):\n",
    "#         return self._make_head()\n",
    "#     @property\n",
    "#     def body(self):\n",
    "#         return self._make_body()\n",
    "    \n",
    "#     def _make_stem(self):\n",
    "#         stem = []\n",
    "#         for i in range(len(self.stem_sizes)-1):\n",
    "#             stem.append((f\"conv_{i}\", self.conv_layer(self.stem_sizes[i], self.stem_sizes[i+1], stride=2 if i==0 else 1)))\n",
    "#         stem.append(('stem_pool', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)))\n",
    "#         return nn.Sequential(OrderedDict(stem))\n",
    "    \n",
    "#     def _make_head(self):\n",
    "#         head = [('pool', nn.AdaptiveAvgPool2d(1)),\n",
    "#                 ('flat', Flatten()),\n",
    "#                 ('fc',   nn.Linear(self.block_szs[-1]*self.expansion, self.c_out))]\n",
    "#         return nn.Sequential(OrderedDict(head))\n",
    "    \n",
    "#     def _make_body(self):\n",
    "#         blocks = [(f\"l_{i}\", self._make_layer(self.expansion, \n",
    "#                         self.block_szs[i], self.block_szs[i+1], l, 1 if i==0 else 2))\n",
    "#                   for i,l in enumerate(self.layers)]\n",
    "#         return nn.Sequential(OrderedDict(blocks))\n",
    "    \n",
    "#     def _make_layer(self, expansion, ni, nf, blocks, stride):\n",
    "#         return nn.Sequential(\n",
    "#             *[self.block(expansion, ni if i==0 else nf, nf, \n",
    "#                          stride if i==0 else 1,conv_layer=self.conv_layer)\n",
    "#               for i in range(blocks)])\n",
    "#     def __call__(self):\n",
    "#         model = nn.Sequential(OrderedDict([\n",
    "# #             ('stem', nn.Sequential(*self.stem)),\n",
    "#             ('stem', self.stem),\n",
    "#             ('body', self.body),\n",
    "#             ('head', self.head)\n",
    "#         ]))\n",
    "#         init_cnn(model)\n",
    "#         return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # v6\n",
    "# class XResNet():\n",
    "#     def __init__(self, expansion=1, layers=[2,2,2,2], c_in=3, c_out=1000):\n",
    "#         super().__init__()\n",
    "#         self.c_in, self.c_out,self.expansion,self.layers = c_in,c_out,expansion,layers\n",
    "#         self.stem_sizes = [c_in,32,32,64]\n",
    "#         self.block_szs = [64//expansion,64,128,256,512] +[256]*(len(layers)-4)\n",
    "#         self.block = ResBlock\n",
    "#         self.norm = nn.BatchNorm2d\n",
    "#         self.act_fn=nn.ReLU(inplace=True)\n",
    "#         self.stem_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "#         self.pool = nn.AvgPool2d(2, ceil_mode=True)\n",
    "#         self.sa=False\n",
    "# #         self._conv_layer = conv_layer\n",
    "    \n",
    "# #     @property\n",
    "# #     def conv_layer(self): return self._conv_layer\n",
    "\n",
    "#     def conv_layer(self, ni, nf, ks=3, stride=1, zero_bn=False, act=True):\n",
    "#         bn = self.norm(nf)\n",
    "#         nn.init.constant_(bn.weight, 0. if zero_bn else 1.)\n",
    "#         layers = [('conv', conv(ni, nf, ks, stride=stride)), ('norm', bn)]\n",
    "#         if act: layers.append(('act_fn', self.act_fn))\n",
    "#         return nn.Sequential(OrderedDict(layers))\n",
    "    \n",
    "#     @property\n",
    "#     def stem(self):\n",
    "#         return self._make_stem()\n",
    "#     @property\n",
    "#     def head(self):\n",
    "#         return self._make_head()\n",
    "#     @property\n",
    "#     def body(self):\n",
    "#         return self._make_body()\n",
    "    \n",
    "#     def _make_stem(self):\n",
    "#         stem = []\n",
    "#         for i in range(len(self.stem_sizes)-1):\n",
    "#             stem.append((f\"conv_{i}\", self.conv_layer(self.stem_sizes[i], self.stem_sizes[i+1], stride=2 if i==0 else 1)))\n",
    "#         stem.append(('stem_pool', self.stem_pool))\n",
    "#         return nn.Sequential(OrderedDict(stem))\n",
    "    \n",
    "#     def _make_head(self):\n",
    "#         head = [('pool', nn.AdaptiveAvgPool2d(1)),\n",
    "#                 ('flat', Flatten()),\n",
    "#                 ('fc',   nn.Linear(self.block_szs[-1]*self.expansion, self.c_out))]\n",
    "#         return nn.Sequential(OrderedDict(head))\n",
    "    \n",
    "#     def _make_body(self):\n",
    "#         blocks = [(f\"l_{i}\", self._make_layer(self.expansion, \n",
    "#                         self.block_szs[i], self.block_szs[i+1], l, \n",
    "#                         1 if i==0 else 2, self.sa if i==0 else False))\n",
    "#                   for i,l in enumerate(self.layers)]\n",
    "#         return nn.Sequential(OrderedDict(blocks))\n",
    "    \n",
    "#     def _make_layer(self, expansion, ni, nf, blocks, stride,sa):\n",
    "#         return nn.Sequential(OrderedDict(\n",
    "#             [(f\"bl_{i}\", self.block(expansion, ni if i==0 else nf, nf, \n",
    "#                     stride if i==0 else 1, sa=sa if i==blocks-1 else False,\n",
    "#                     conv_layer=self.conv_layer, act_fn=self.act_fn, pool=self.pool))\n",
    "#               for i in range(blocks)]))\n",
    "#     def __call__(self):\n",
    "#         model = nn.Sequential(OrderedDict([\n",
    "# #             ('stem', nn.Sequential(*self.stem)),\n",
    "#             ('stem', self.stem),\n",
    "#             ('body', self.body),\n",
    "#             ('head', self.head)\n",
    "#         ]))\n",
    "#         init_cnn(model)\n",
    "#         return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# XResNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # export\n",
    "# # v6\n",
    "# class XResNet():\n",
    "#     def __init__(self, expansion=1, layers=[2,2,2,2], c_in=3, c_out=1000):\n",
    "#         super().__init__()\n",
    "#         self.c_in, self.c_out,self.expansion,self.layers = c_in,c_out,expansion,layers\n",
    "#         self.stem_sizes = [c_in,32,32,64]\n",
    "#         self.stem_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "#         self.stem_bn_end = False\n",
    "#         self.block_szs = [64//expansion,64,128,256,512] +[256]*(len(layers)-4)\n",
    "#         self.block = ResBlock\n",
    "#         self.norm = nn.BatchNorm2d\n",
    "#         self.act_fn=nn.ReLU(inplace=True)\n",
    "#         self.pool = nn.AvgPool2d(2, ceil_mode=True)\n",
    "#         self.sa=False\n",
    "#         self.bn_1st = True\n",
    "#         #         self._conv_layer = conv_layer\n",
    "#         #     @property\n",
    "#         #     def conv_layer(self): return self._conv_layer\n",
    "\n",
    "#     def conv_layer(self, ni, nf, ks=3, stride=1, zero_bn=False, act=True, norm=True):\n",
    "#         bn = self.norm(nf)\n",
    "#         nn.init.constant_(bn.weight, 0. if zero_bn else 1.)\n",
    "#         layers =  [('norm', bn)] if norm else []\n",
    "#         if act: layers.append(('act_fn', self.act_fn))\n",
    "#         if not self.bn_1st: layers.reverse()\n",
    "#         layers = [('conv', conv(ni, nf, ks, stride=stride))] + layers\n",
    "#         return nn.Sequential(OrderedDict(layers))\n",
    "    \n",
    "#     @property\n",
    "#     def stem(self):\n",
    "#         return self._make_stem()\n",
    "#     @property\n",
    "#     def head(self):\n",
    "#         return self._make_head()\n",
    "#     @property\n",
    "#     def body(self):\n",
    "#         return self._make_body()\n",
    "    \n",
    "#     def _make_stem(self):\n",
    "#         stem = [(f\"conv_{i}\", self.conv_layer(self.stem_sizes[i], self.stem_sizes[i+1], \n",
    "#                         stride=2 if i==0 else 1, \n",
    "#                         norm=(not self.stem_bn_end) if i==(len(self.stem_sizes)-2) else True))\n",
    "#                 for i in range(len(self.stem_sizes)-1)]\n",
    "#         stem.append(('stem_pool', self.stem_pool))\n",
    "#         if self.stem_bn_end: stem.append(('norm', self.norm(self.stem_sizes[-1])))\n",
    "#         return nn.Sequential(OrderedDict(stem))\n",
    "    \n",
    "#     def _make_head(self):\n",
    "#         head = [('pool', nn.AdaptiveAvgPool2d(1)),\n",
    "#                 ('flat', Flatten()),\n",
    "#                 ('fc',   nn.Linear(self.block_szs[-1]*self.expansion, self.c_out))]\n",
    "#         return nn.Sequential(OrderedDict(head))\n",
    "    \n",
    "#     def _make_body(self):\n",
    "#         blocks = [(f\"l_{i}\", self._make_layer(self.expansion, \n",
    "#                         self.block_szs[i], self.block_szs[i+1], l, \n",
    "#                         1 if i==0 else 2, self.sa if i==0 else False))\n",
    "#                   for i,l in enumerate(self.layers)]\n",
    "#         return nn.Sequential(OrderedDict(blocks))\n",
    "    \n",
    "#     def _make_layer(self, expansion, ni, nf, blocks, stride,sa):\n",
    "#         return nn.Sequential(OrderedDict(\n",
    "#             [(f\"bl_{i}\", self.block(expansion, ni if i==0 else nf, nf, \n",
    "#                     stride if i==0 else 1, sa=sa if i==blocks-1 else False,\n",
    "#                     conv_layer=self.conv_layer, act_fn=self.act_fn, pool=self.pool))\n",
    "#               for i in range(blocks)]))\n",
    "    \n",
    "#     def __call__(self):\n",
    "#         model = nn.Sequential(OrderedDict([\n",
    "#             ('stem', self.stem),\n",
    "#             ('body', self.body),\n",
    "#             ('head', self.head)\n",
    "#         ]))\n",
    "#         init_cnn(model)\n",
    "#         return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # v7\n",
    "# class Net():\n",
    "#     def __init__(self, expansion=1, layers=[2,2,2,2], c_in=3, c_out=1000, name='Net'):\n",
    "#         super().__init__()\n",
    "#         self.name = name\n",
    "#         self.c_in, self.c_out,self.expansion,self.layers = c_in,c_out,expansion,layers\n",
    "#         self.stem_sizes = [c_in,32,32,64]\n",
    "#         self.stem_pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "#         self.stem_bn_end = False\n",
    "#         self.block_szs = [64//expansion,64,128,256,512] +[256]*(len(layers)-4)\n",
    "#         self.block = ResBlock\n",
    "#         self.norm = nn.BatchNorm2d\n",
    "#         self.act_fn=nn.ReLU(inplace=True)\n",
    "#         self.pool = nn.AvgPool2d(2, ceil_mode=True)\n",
    "#         self.sa=False\n",
    "#         self.bn_1st = True\n",
    "#         self.zero_bn=True\n",
    "\n",
    "#     def conv_layer(self, ni, nf, ks=3, stride=1, zero_bn=False, act=True, norm=True):\n",
    "#         bn = self.norm(nf)\n",
    "#         nn.init.constant_(bn.weight, 0. if zero_bn else 1.)\n",
    "#         layers =  [('norm', bn)] if norm else []\n",
    "#         if act: layers.append(('act_fn', self.act_fn))\n",
    "#         if not self.bn_1st: layers.reverse()\n",
    "#         layers = [('conv', conv(ni, nf, ks, stride=stride))] + layers\n",
    "#         return nn.Sequential(OrderedDict(layers))\n",
    "    \n",
    "#     @property\n",
    "#     def stem(self):\n",
    "#         return self._make_stem()\n",
    "#     @property\n",
    "#     def head(self):\n",
    "#         return self._make_head()\n",
    "#     @property\n",
    "#     def body(self):\n",
    "#         return self._make_body()\n",
    "    \n",
    "#     def _make_stem(self):\n",
    "#         stem = [(f\"conv_{i}\", self.conv_layer(self.stem_sizes[i], self.stem_sizes[i+1], \n",
    "#                         stride=2 if i==0 else 1, \n",
    "#                         norm=(not self.stem_bn_end) if i==(len(self.stem_sizes)-2) else True))\n",
    "#                 for i in range(len(self.stem_sizes)-1)]\n",
    "#         stem.append(('stem_pool', self.stem_pool))\n",
    "#         if self.stem_bn_end: stem.append(('norm', self.norm(self.stem_sizes[-1])))\n",
    "#         return nn.Sequential(OrderedDict(stem))\n",
    "    \n",
    "#     def _make_head(self):\n",
    "#         head = [('pool', nn.AdaptiveAvgPool2d(1)),\n",
    "#                 ('flat', Flatten()),\n",
    "#                 ('fc',   nn.Linear(self.block_szs[-1]*self.expansion, self.c_out))]\n",
    "#         return nn.Sequential(OrderedDict(head))\n",
    "    \n",
    "#     def _make_body(self):\n",
    "#         blocks = [(f\"l_{i}\", self._make_layer(self.expansion, \n",
    "#                         self.block_szs[i], self.block_szs[i+1], l, \n",
    "#                         1 if i==0 else 2, self.sa if i==0 else False))\n",
    "#                   for i,l in enumerate(self.layers)]\n",
    "#         return nn.Sequential(OrderedDict(blocks))\n",
    "    \n",
    "#     def _make_layer(self,expansion,ni,nf,blocks,stride,sa):\n",
    "#         return nn.Sequential(OrderedDict(\n",
    "#             [(f\"bl_{i}\", self.block(expansion, ni if i==0 else nf, nf, \n",
    "#                     stride if i==0 else 1, sa=sa if i==blocks-1 else False,\n",
    "#                     conv_layer=self.conv_layer, act_fn=self.act_fn, pool=self.pool,zero_bn=self.zero_bn))\n",
    "#               for i in range(blocks)]))\n",
    "    \n",
    "#     def __call__(self):\n",
    "#         model = nn.Sequential(OrderedDict([\n",
    "#             ('stem', self.stem),\n",
    "#             ('body', self.body),\n",
    "#             ('head', self.head)\n",
    "#         ]))\n",
    "#         init_cnn(model)\n",
    "#         model.extra_repr = lambda : f\"model {self.name}\"\n",
    "#         return model\n",
    "#     def __repr__(self):\n",
    "#         return f\" constr {self.name}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model  = XResNet()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.stem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.bn_1st = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.act_fn =nn.LeakyReLU(inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.sa = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.stem_bn_end = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# m = model()\n",
    "# m"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# bs_test = 16\n",
    "# xb = torch.randn(bs_test, 3, 128, 128)\n",
    "# y = m(xb)\n",
    "# # y.shape\n",
    "# print(y.shape)\n",
    "# assert y.shape == torch.Size([bs_test, 1000]), f\"size expected {bs_test}, 1000\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# m.stem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# m.body.l_0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.conv_layer(64,64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# model.stem_sizes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## xresnet constructor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # 1 ver \n",
    "# def xresnet(expansion, n_layers, name, c_out=1000, pretrained=False, **kwargs):\n",
    "#     model = XResNet(expansion, n_layers, c_out=c_out, **kwargs)\n",
    "# #     return model\n",
    "#     return model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # v1\n",
    "# me = sys.modules[__name__]\n",
    "# for n,e,l in [\n",
    "#     [ 18 , 1, [2,2,2 ,2] ],\n",
    "#     [ 34 , 1, [3,4,6 ,3] ],\n",
    "#     [ 50 , 4, [3,4,6 ,3] ],\n",
    "#     [ 101, 4, [3,4,23,3] ],\n",
    "#     [ 152, 4, [3,8,36,3] ],\n",
    "# ]:\n",
    "#     name = f'xresnet{n}'\n",
    "#     setattr(me, name, partial(xresnet, expansion=e, n_layers=l, name=name))\n",
    "\n",
    "# xresnet18_deep = partial(xresnet, expansion=1, n_layers=[2, 2,  2, 2,1,1], name='xresnet18_deep')\n",
    "# xresnet34_deep = partial(xresnet, expansion=1, n_layers=[3, 4,  6, 3,1,1], name='xresnet34_deep')\n",
    "# xresnet50_deep = partial(xresnet, expansion=4, n_layers=[3, 4,  6, 3,1,1], name='xresnet50_deep')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# def xresnet(expansion, n_layers, name, c_out=1000, pretrained=False, **kwargs):\n",
    "#     model = XResNet(expansion, n_layers, c_out=c_out, **kwargs)\n",
    "#     return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # export\n",
    "# def xresnet(expansion, n_layers, name, c_out=1000, pretrained=False, **kwargs):\n",
    "#     return partial(XResNet, expansion, n_layers, c_out=c_out, **kwargs)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# xresnet50_d = xresnet(expansion=4, n_layers=[3, 4,  6, 3,1,1], name='xresnet50_deep')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# m = xresnet50_d(c_out=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# m.c_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # export\n",
    "# me = sys.modules[__name__]\n",
    "# for n,e,l in [\n",
    "#     [ 18 , 1, [2,2,2 ,2] ],\n",
    "#     [ 34 , 1, [3,4,6 ,3] ],\n",
    "#     [ 50 , 4, [3,4,6 ,3] ],\n",
    "#     [ 101, 4, [3,4,23,3] ],\n",
    "#     [ 152, 4, [3,8,36,3] ],\n",
    "# ]:\n",
    "#     name = f'xresnet{n}'\n",
    "#     setattr(me, name, xresnet(expansion=e, n_layers=l, name=name))\n",
    "\n",
    "# xresnet18_deep = xresnet(expansion=1, n_layers=[2, 2,  2, 2,1,1], name='xresnet18_deep')\n",
    "# xresnet34_deep = xresnet(expansion=1, n_layers=[3, 4,  6, 3,1,1], name='xresnet34_deep')\n",
    "# xresnet50_deep = xresnet(expansion=4, n_layers=[3, 4,  6, 3,1,1], name='xresnet50_deep')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# me = sys.modules[__name__]\n",
    "# for n,e,l in [[ 18 , 1, [2,2,2 ,2] ],\n",
    "#     [ 34 , 1, [3,4,6 ,3] ],\n",
    "#     [ 50 , 4, [3,4,6 ,3] ],\n",
    "#     [ 101, 4, [3,4,23,3] ],\n",
    "#     [ 152, 4, [3,8,36,3] ],]:\n",
    "#     name = f'net{n}'\n",
    "#     setattr(me, name, partial(XResNet, expansion=e, layers=l,))# name=name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# end\n",
    "model_constructor\n",
    "by ayasyrev"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # hide\n",
    "# from nbdev.export import *\n",
    "# notebook2script()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
